#ifndef BOOTSTRAPPING_BOOTSTRAPPING_H
#define BOOTSTRAPPING_BOOTSTRAPPING_H
#include <chrono>
#include "transform.h"
#include "contextchain.h"

/* *
 * @class SlotwiseNorm
 * @brief 封装每一个槽中norm的计算
 */
class SlotwiseNorm {
    const SlotRing &slotRing;
    size_t subfieldIndexLog2;

public:
    SlotwiseNorm(const SlotRing &slotRing, size_t subfieldIndexLog2);
    SlotwiseNorm(const SlotwiseNorm &) = default;
    SlotwiseNorm(SlotwiseNorm &&) = default;
    ~SlotwiseNorm() = default;

    poly operator () (const poly &x) const;
    void ApplyCiphertext(Ciphertext &in, Evaluator &eval, GaloisKeys &gk, Ciphertext &destination) const;
    std::vector<uint32_t> GaloisElements() const;
};

/* *
 * @class SlotwiseTrace
 * @brief 封装每一个槽中trace的计算
 */
class SlotwiseTrace {
    const SlotRing &slotRing;
    size_t sourceSubfieldIndexLog2;
    size_t targetSubfieldIndexLog2;

public:
    SlotwiseTrace(const SlotRing &slotRing, size_t sourceSubfieldIndexLog2, size_t targetSubfieldIndexLog2);
    SlotwiseTrace(const SlotwiseTrace &) = default;
    SlotwiseTrace(SlotwiseTrace &&) = default;
    ~SlotwiseTrace() = default;

    poly operator () (const poly &x) const;
    void ApplyCiphertext(Ciphertext &in, Evaluator &eval, GaloisKeys &gk, Ciphertext &destination) const;
    std::vector<uint32_t> GaloisElements() const;
    size_t field_index() const;
};

/* *
 * @class PolyEvaluator
 * @brief 所有可以评估数字提取的对象的抽象基类。
 */
class PolyEvaluator {
public:
    PolyEvaluator() = default;
    PolyEvaluator(const PolyEvaluator &) = delete;
    PolyEvaluator(PolyEvaluator &&) = delete;
    virtual ~PolyEvaluator() = default;

    virtual poly operator () (const poly &) const = 0;
    virtual void ApplyCiphertext(Ciphertext &in, Context &context, Evaluator &eval, GaloisKeys &gk,
        Ciphertext &dst) const = 0;
    virtual std::vector<uint32_t> GaloisElements() const = 0;
};

/* *
 * @class P127PolyEvaluator
 * @brief 该类用于实现同态加密中的P127多项式评估。
 *
 * `P127PolyEvaluator` 类负责对同态加密中的多项式进行评估，并应用修正多项式。
 */
class P127PolyEvaluator : public PolyEvaluator {
    const SlotRing &slotRing;
    poly evaluationElement;
    size_t log2Exponent;
    SlotwiseNorm normOp;

    poly correctionPoly;

public:
    P127PolyEvaluator(const SlotRing &slotRing, poly evaluationElement_, poly correctionPoly_, size_t log2Exponent_);
    virtual ~P127PolyEvaluator() = default;

    virtual poly operator () (const poly &x) const;
    virtual void ApplyCiphertext(Ciphertext &in, Context &context, Evaluator &eval, GaloisKeys &gk,
        Ciphertext &destination) const;
    virtual std::vector<uint32_t> GaloisElements() const;
};

/* *
 * @class P257CorrectionPolyEvaluator
 * @brief 该类用于实现同态加密中的P257修正多项式评估。
 */
class P257CorrectionPolyEvaluator : public PolyEvaluator {
    const SlotRing &slotRing;
    poly evaluationElement;
    size_t log2Exponent;
    SlotwiseNorm normOp;
    poly correctionPoly;
    uint64_t constantCorrection;

public:
    P257CorrectionPolyEvaluator(const SlotRing &slotRing, poly evaluationElement, poly correctionPoly,
        size_t log2Exponent);
    virtual ~P257CorrectionPolyEvaluator() = default;

    virtual poly operator () (const poly &x) const;
    virtual void ApplyCiphertext(Ciphertext &in, Context &context, Evaluator &eval, GaloisKeys &gk,
        Ciphertext &destination) const;
    virtual std::vector<uint32_t> GaloisElements() const;
};

class Bootstrapper;
/* *
 * @class BootstrappingKey
 * @brief 该类表示同态加密中的自举密钥（Bootstrapping Key）。
 *
 * `BootstrappingKey` 类用于存储和管理同态加密中的自举操作所需要的密钥。自举密钥由以下几个部分组成：
 * - 加密后的私钥（`encryptedSk`）
 * - 伽罗瓦密钥（`GaloisKeys`）
 * - 重线性化密钥（`RelinearizationKey`）
 *
 * 自举密钥在同态加密中用于支持自举操作，减少密文计算过程中的噪声积累，并使得加密计算能够在多个计算阶段之间进行。
 */
class BootstrappingKey {
    Ciphertext encryptedSk;
    GaloisKeys gk;
    RelinearizationKey rk;

public:
    BootstrappingKey() = default;
    BootstrappingKey(const BootstrappingKey &) = default;
    BootstrappingKey(BootstrappingKey &&) = default;
    ~BootstrappingKey() = default;

    const GaloisKeys &GalKeys() const;
    const RelinearizationKey &RelinKeys() const;

    friend Bootstrapper;
};

// 在p=127下硬解码数字提取多项式测试参数
std::unique_ptr<PolyEvaluator> P127TestParametersDigitExtractor(const SlotRing &slotRing);

// 在p=257下硬解码数字提取多项式测试参数
std::unique_ptr<PolyEvaluator> P257TestParametersDigitExtractor(const SlotRing &slotRing);

void TestEncodeDecode();
void TestGaloisPolyEvaluator();
void TestHomomorphicNoisyDecrypt();
void TestSlotwiseDigitExtract();
void TestCoeffsToSlots();
void TestBootstrapping();

/* *
 * @class Bootstrapper
 * @brief 该类实现了同态加密中的自举操作（Bootstrapping）。
 *
 * `Bootstrapper` 类提供了同态加密中自举操作的实现，包括噪声解密、槽位数字提取、系数与槽位之间的转换等。
 * 通过这些操作，可以在密文计算过程中减少噪声的积累，支持更深的加密计算。该类包含了对自举密钥生成、初始化及各种自举操作的支持。
 *
 * 主要功能包括：
 * - 同态解密操作（Homomorphic Noisy Decrypt）
 * - 槽位数字提取（Slotwise Digit Extract）
 * - 槽位与系数的转换操作（Slots to Coeffs，Coeffs to Slots）
 * - 线性变换初始化和自举密钥的创建
 */
class Bootstrapper {
public:
    ContextChain contextChain;
    std::shared_ptr<const SlotRing> slotRing;

    std::unique_ptr<PolyEvaluator> digitExtractPoly = nullptr;
    std::unique_ptr<SlotwiseTrace> traceOp = nullptr;
    std::unique_ptr<CompiledSubringLinearTransform> slotsToCoefficients = nullptr;
    std::unique_ptr<CompiledSubringLinearTransform> coefficientsToSlots = nullptr;

    size_t PolyModulusDegree() const noexcept;
    Evaluator &BootstrappingEvaluator() const;
    const Context &BootstrappingContext() const;

public:

    void Bootstrapping(Ciphertext &ciphertext, SecretKey &sk, Context &context);
    
    void HomomorphicNoisyDecrypt(Ciphertext &ciphertext, BootstrappingKey &bk, Ciphertext &destination) const;

    void SlotwiseDigitExtract(Ciphertext &ciphertext, BootstrappingKey &bk, Ciphertext &destination) const;

    void SlotsToCoeffs(Ciphertext &ciphertext, BootstrappingKey &bk, Ciphertext &destination) const;

    void CoeffsToSlots(Ciphertext &ciphertext, BootstrappingKey &bk, Ciphertext &destination) const;

public:
    Bootstrapper(const Context &bootstrappedContext, std::shared_ptr<const SlotRing> slotRing,
        std::unique_ptr<PolyEvaluator> digit_extract_poly);
    Bootstrapper(const Bootstrapper &) = default;
    Bootstrapper(Bootstrapper &&) = default;
    ~Bootstrapper() = default;

    // 初始化线性变换参数
    void Initialize();
    // 创建自举中伽罗瓦密钥和重线性化密钥
    void CreateBootstrappingKey(SecretKey &sk, BootstrappingKey &bk, bool allowUninitializedForTest = false) const;
    // 根据初始密钥生成自举密钥
    void CreateSecretKey(const SecretKey &base_sk, SecretKey &destination) const;

    friend void TestHomomorphicNoisyDecrypt();
    friend void TestSlotwiseDigitExtract();
};

inline size_t Bootstrapper::PolyModulusDegree() const noexcept
{
    return contextChain.PolyModulusDegree();
}

inline Evaluator &Bootstrapper::BootstrappingEvaluator() const
{
    return contextChain.TargetEvaluator();
}

inline const GaloisKeys &BootstrappingKey::GalKeys() const
{
    return gk;
}
inline const RelinearizationKey &BootstrappingKey::RelinKeys() const
{
    return rk;
}
#endif